import torch
import torch.nn as nn
import torch.nn.functional as F

# 基本卷积块，根据参数选择构造2d或3d、卷积或转置卷积、批归一化和ReLU
class BasicConv(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, **kwargs):
        super(BasicConv, self).__init__()

        self.relu = relu # 设置是否使用ReLU(LeakyReLU)
        self.use_bn = bn # 设置是否使用批归一化

        if is_3d:
            if deconv:
                self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs)
            else:
                self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
            self.bn = nn.BatchNorm3d(out_channels)
        else:
            if deconv:
                self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs)
            else:
                self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
            self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        x = self.conv(x)
        if self.use_bn:
            x = self.bn(x)
        if self.relu:
            x = nn.LeakyReLU()(x) # 使用LeakyReLU，防止ReLU死亡
        return x

# 残差块，拼接张量或残差连接
class Conv2x(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, relu=True, keep_dispc=False):
        super(Conv2x, self).__init__()

        self.concat = concat # 是否拼接张量
        self.is_3d = is_3d
        if deconv and is_3d:
            kernel = (4, 4, 4) # 3d转置卷积核
        elif deconv:
            kernel = 4 # 2d转置卷积核
        else:
            kernel = 3

        if deconv and is_3d and keep_dispc:
            kernel = (1, 4, 4)
            stride = (1, 2, 2)
            padding = (0, 1, 1)
            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, stride=stride, padding=padding) # 3d转置卷积，LeakyReLU，深度不变，宽高*2
        else:
            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, stride=2, padding=1) # LeakyReLU，尺寸都/2或*2

        if self.concat:
            mul = 2 if keep_concat else 1
            self.conv2 = BasicConv(out_channels*2, out_channels*mul, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) # 卷积，输出通道*2，输出通道*2(或不变)，尺寸不变
        else:
            self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1) # 卷积，尺寸不变

    def forward(self, x, rem):
        x = self.conv1(x)
        if x.shape != rem.shape: # 连接的两层张量形状不同
            x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='nearest') # 最近邻插值x为rem形状
        if self.concat:
            x = torch.cat((x, rem), 1) # 沿着通道维度拼接张量，即通道*2
        else:
            x = x + rem # 残差连接
        x = self.conv2(x) # 卷积，尺寸不变
        return x

# 基本卷积块，带有实例归一化
class BasicConv_IN(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs):
        super(BasicConv_IN, self).__init__()

        self.relu = relu
        self.use_in = IN # 是否使用实例归一化
        if is_3d:
            if deconv:
                self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) # 3D转置卷积
            else:
                self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
            self.IN = nn.InstanceNorm3d(out_channels) # 3d实例归一化
        else:
            if deconv:
                self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) # 2D转置卷积
            else:
                self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
            self.IN = nn.InstanceNorm2d(out_channels) # 2d实例归一化

    def forward(self, x):
        x = self.conv(x)
        if self.use_in:
            x = self.IN(x)
        if self.relu:
            x = nn.LeakyReLU()(x) # 使用LeakyReLU，防止ReLU死亡
        return x

# 残差块，拼接张量或残差连接，带有实例归一化
class Conv2x_IN(nn.Module):
    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, relu=True, keep_dispc=False):
        super(Conv2x_IN, self).__init__()

        self.concat = concat # 是否拼接张量
        self.is_3d = is_3d
        if deconv and is_3d:
            kernel = (4, 4, 4) # 3d转置卷积核
        elif deconv:
            kernel = 4 # 2d转置卷积核
        else:
            kernel = 3

        if deconv and is_3d and keep_dispc:
            kernel = (1, 4, 4)
            stride = (1, 2, 2)
            padding = (0, 1, 1)
            self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=stride, padding=padding) # 3d转置卷积，实例归一化，LeakyReLU，深度不变，宽高*2
        else:
            self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=2, padding=1) # 实例归一化，LeakyReLU，尺寸都/2或*2

        if self.concat:
            mul = 2 if keep_concat else 1
            self.conv2 = BasicConv_IN(out_channels*2, out_channels*mul, False, is_3d, IN, relu, kernel_size=3, stride=1, padding=1) # 卷积，输出通道*2，输出通道*2(或不变)，尺寸不变
        else:
            self.conv2 = BasicConv_IN(out_channels, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, padding=1) # 卷积，尺寸不变

    def forward(self, x, rem):
        x = self.conv1(x)
        if x.shape != rem.shape: # 连接的两层张量形状不同
            x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='nearest') # 最近邻插值x为rem形状
        if self.concat:
            x = torch.cat((x, rem), 1) # 沿着通道维度拼接张量，即通道*2
        else:
            x = x + rem # 残差连接
        x = self.conv2(x) # 卷积，尺寸不变
        return x

# 计算组相关性
def groupwise_correlation(fea1, fea2, num_groups):
    B, C, H, W = fea1.shape
    assert C % num_groups == 0
    channels_per_group = C // num_groups
    cost = (fea1 * fea2).view([B, num_groups, channels_per_group, H, W]).mean(dim=2) # 计算张量相似度，张量相乘，(B,组数,每组通道数,H,W)，再对每组通道数维度取平均值
    assert cost.shape == (B, num_groups, H, W)
    return cost # 输出(B,组数,H,W)

# 构造组相关性成本体积
def build_gwc_volume(refimg_fea, targetimg_fea, maxdisp, num_groups):
    B, C, H, W = refimg_fea.shape
    volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W]) # 创建零张量，(B,组数,最大视差,H,W)
    for i in range(maxdisp):
        if i > 0:
            volume[:, :, i, :, i:] = groupwise_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], num_groups) # 将左视图切片i视差，右视图右移i视差，计算两个张量的组相关性，输出(B,组数,视差,H,W)，宽不够的为0
        else:
            volume[:, :, i, :, :] = groupwise_correlation(refimg_fea, targetimg_fea, num_groups) # 0视差时，计算不偏移的左视图和右视图的相关性
    volume = volume.contiguous()
    return volume # 输出(B,组数,视差,H,W)

def norm_correlation(fea1, fea2):
    cost = torch.mean(((fea1/(torch.norm(fea1, 2, 1, True)+1e-05)) * (fea2/(torch.norm(fea2, 2, 1, True)+1e-05))), dim=1, keepdim=True)
    return cost

def build_norm_correlation_volume(refimg_fea, targetimg_fea, maxdisp):
    B, C, H, W = refimg_fea.shape
    volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W])
    for i in range(maxdisp):
        if i > 0:
            volume[:, :, i, :, i:] = norm_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i])
        else:
            volume[:, :, i, :, :] = norm_correlation(refimg_fea, targetimg_fea)
    volume = volume.contiguous()
    return volume

def correlation(fea1, fea2):
    cost = torch.sum((fea1 * fea2), dim=1, keepdim=True)
    return cost

def build_correlation_volume(refimg_fea, targetimg_fea, maxdisp):
    B, C, H, W = refimg_fea.shape
    volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W])
    for i in range(maxdisp):
        if i > 0:
            volume[:, :, i, :, i:] = correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i])
        else:
            volume[:, :, i, :, :] = correlation(refimg_fea, targetimg_fea)
    volume = volume.contiguous()
    return volume

def build_concat_volume(refimg_fea, targetimg_fea, maxdisp):
    B, C, H, W = refimg_fea.shape
    volume = refimg_fea.new_zeros([B, 2 * C, maxdisp, H, W])
    for i in range(maxdisp):
        if i > 0:
            volume[:, :C, i, :, :] = refimg_fea[:, :, :, :]
            volume[:, C:, i, :, i:] = targetimg_fea[:, :, :, :-i]
        else:
            volume[:, :C, i, :, :] = refimg_fea
            volume[:, C:, i, :, :] = targetimg_fea
    volume = volume.contiguous()
    return volume

# 回归视差，输入概率体积，当前尺寸最大视差
def disparity_regression(x, maxdisp):
    assert len(x.shape) == 4 # 概率体积(batch_size,1/4最大视差概率,1/4原高,1/4原宽)
    disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device) # 生成0~(maxdisp-1)的张量[0, 1, ...]
    disp_values = disp_values.view(1, maxdisp, 1, 1) # 调整张量形状为(1,maxdisp,1,1)
    return torch.sum(x * disp_values, 1, keepdim=True) # 计算x*disp_values(概率维度*视差维度)，在第1维度(*的维度)求和，保持维度数量一致，返回(batch_size,1,1/4原高,1/4原宽)

# CoEx相关和激励
class FeatureAtt(nn.Module):
    def __init__(self, cv_chan, feat_chan):
        super(FeatureAtt, self).__init__()

        self.feat_att = nn.Sequential(
            BasicConv(feat_chan, feat_chan//2, kernel_size=1, stride=1, padding=0), # 2d卷积，批归一化，LeakyReLU，输出为输入通道//2，尺寸不变
            nn.Conv2d(feat_chan//2, cv_chan, 1)) # 2d卷积，尺寸不变

    def forward(self, cv, feat):
        feat_att = self.feat_att(feat).unsqueeze(2) # 经过2次2d卷积，尺寸不变，在第2个维度上添加一个维度，输出(batch_size,cv_chan,添加维度,高,宽)
        cv = torch.sigmoid(feat_att)*cv # CoEx相关和激励，sigmoid激活特征后*cv
        return cv

# 上下文上采样
def context_upsample(disp_low, up_weights):
    ###
    # cv (b,1,h,w)
    # sp (b,9,4*h,4*w)
    ###
    b, c, h, w = disp_low.shape # 1/4原尺寸视差

    disp_unfold = F.unfold(disp_low.reshape(b,c,h,w), 3, 1, 1).reshape(b,-1,h,w) # 展开张量，滑窗3*3，步长1，每次覆盖形成1个列向量，输出(B,C,9,H*W)，改变张量形状为(B,C*9,H,W)
    disp_unfold = F.interpolate(disp_unfold, (h*4,w*4), mode='nearest').reshape(b,9,h*4,w*4) # 粗分辨率视差，最近邻插值disp_unfold为(H*4,W*4)的图像形状，改变张量形状为(B,9,H*4,W*4)，即原尺寸

    disp = (disp_unfold*up_weights).sum(1) # 上采样视差，粗分辨率视差*上下文上采样权重，在通道维度求和

    return disp # 返回原尺寸视差，(B,1,H,W)
